from mmseg.apis import inference_model as seg_inference_model
from mmseg.apis import init_model as seg_init_model
import mmcv
import sys
import cv2
import numpy as np
from tqdm import tqdm
from glob import glob
import os.path as osp
# from collections import Counter
base_dir = osp.dirname(osp.dirname(osp.dirname(osp.dirname(osp.abspath(__file__)))))
sys.path.append(base_dir)
from ltcv.models.necks import FPNX
from ltcv.models.heads import FPNHeadX

def seg_init(config_path_seg,checkpoint_path_seg):
    model = seg_init_model(config_path_seg, checkpoint_path_seg)
    return model

def seg_infer(model,img_dirs):
    
    results = []
    # img_dirs = glob(osp.join(data_dir, "*.jpeg")) + glob(osp.join(data_dir, "*.jpg")) + glob(osp.join(data_dir, "*.png"))
    for img_dir in tqdm(img_dirs):
        result = seg_inference_model(model, img_dir)
        # mask = result.pred_sem_seg.data.cpu().numpy().transpose(1, 2, 0)
        mask = result.pred_sem_seg.data.cpu().numpy()
        mask = np.squeeze(mask)
        results.append((img_dir, mask))
    return results
